// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// Buffered stream implementations

#pragma once

#include <cstdint>
#include <memory>
#include <string_view>

#include <nebula/io/concurrency.h>
#include <nebula/io/interfaces.h>
#include <turbo/base/macros.h>

namespace nebula {

    class Buffer;

    class MemoryPool;

}
namespace nebula::io {

    class TURBO_EXPORT BufferedOutputStream : public OutputStream {
    public:
        ~BufferedOutputStream() override;

        /// \brief create a buffered output stream wrapping the given output stream.
        /// \param[in] buffer_size the size of the temporary write buffer
        /// \param[in] pool a MemoryPool to use for allocations
        /// \param[in] raw another OutputStream
        /// \return the created BufferedOutputStream
        static turbo::Result<std::shared_ptr<BufferedOutputStream>> create(
                int64_t buffer_size, MemoryPool *pool, std::shared_ptr<OutputStream> raw);

        /// \brief resize internal buffer
        /// \param[in] new_buffer_size the new buffer size
        /// \return turbo::Status
        turbo::Status set_buffer_size(int64_t new_buffer_size);

        /// \brief Return the current size of the internal buffer
        int64_t buffer_size() const;

        /// \brief Return the number of remaining bytes that have not been flushed to
        /// the raw OutputStream
        int64_t bytes_buffered() const;

        /// \brief flush any buffered writes and release the raw
        /// OutputStream. Further operations on this object are invalid
        /// \return the underlying OutputStream
        turbo::Result<std::shared_ptr<OutputStream>> detach();

        // OutputStream interface

        /// \brief close the buffered output stream.  This implicitly closes the
        /// underlying raw output stream.
        turbo::Status close() override;

        turbo::Status abort() override;

        bool closed() const override;

        turbo::Result<int64_t> tell() const override;

        // write bytes to the stream. Thread-safe
        turbo::Status write(const void *data, int64_t nbytes) override;

        turbo::Status write(const std::shared_ptr<Buffer> &data) override;

        turbo::Status flush() override;

        /// \brief Return the underlying raw output stream.
        std::shared_ptr<OutputStream> raw() const;

    private:
        explicit BufferedOutputStream(std::shared_ptr<OutputStream> raw, MemoryPool *pool);

        class TURBO_NO_EXPORT Impl;

        std::unique_ptr<Impl> impl_;
    };

/// \class BufferedInputStream
/// \brief An InputStream that performs buffered reads from an unbuffered
/// InputStream, which can mitigate the overhead of many small reads in some
/// cases
    class TURBO_EXPORT BufferedInputStream
            : public internal::InputStreamConcurrencyWrapper<BufferedInputStream> {
    public:
        ~BufferedInputStream() override;

        /// \brief create a BufferedInputStream from a raw InputStream
        /// \param[in] buffer_size the size of the temporary read buffer
        /// \param[in] pool a MemoryPool to use for allocations
        /// \param[in] raw a raw InputStream
        /// \param[in] raw_read_bound a bound on the maximum number of bytes
        /// to read from the raw input stream. The default -1 indicates that
        /// it is unbounded
        /// \return the created BufferedInputStream
        static turbo::Result<std::shared_ptr<BufferedInputStream>> create(
                int64_t buffer_size, MemoryPool *pool, std::shared_ptr<InputStream> raw,
                int64_t raw_read_bound = -1);

        /// \brief resize internal read buffer; calls to read(...) will read at least
        /// \param[in] new_buffer_size the new read buffer size
        /// \return turbo::Status
        turbo::Status set_buffer_size(int64_t new_buffer_size);

        /// \brief Return the number of remaining bytes in the read buffer
        int64_t bytes_buffered() const;

        /// \brief Return the current size of the internal buffer
        int64_t buffer_size() const;

        /// \brief Release the raw InputStream. Any data buffered will be
        /// discarded. Further operations on this object are invalid
        /// \return raw the underlying InputStream
        std::shared_ptr<InputStream> detach();

        /// \brief Return the unbuffered InputStream
        std::shared_ptr<InputStream> raw() const;

        // InputStream APIs

        bool closed() const override;

        turbo::Result<std::shared_ptr<const KeyValueMetadata>> read_metadata() override;

        Future<std::shared_ptr<const KeyValueMetadata>> read_metadata_async(
                const IOContext &io_context) override;

    private:
        friend InputStreamConcurrencyWrapper<BufferedInputStream>;

        explicit BufferedInputStream(std::shared_ptr<InputStream> raw, MemoryPool *pool,
                                     int64_t raw_total_bytes_bound);

        turbo::Status DoClose();

        turbo::Status DoAbort() override;

        /// \brief Returns the position of the buffered stream, though the position
        /// of the unbuffered stream may be further advanced.
        turbo::Result<int64_t> DoTell() const;

        turbo::Result<int64_t> DoRead(int64_t nbytes, void *out);

        /// \brief read into buffer.
        turbo::Result<std::shared_ptr<Buffer>> DoRead(int64_t nbytes);

        /// \brief Return a zero-copy string view referencing buffered data,
        /// but do not advance the position of the stream. Buffers data and
        /// expands the buffer size if necessary
        turbo::Result<std::string_view> DoPeek(int64_t nbytes) override;

        class TURBO_NO_EXPORT Impl;

        std::unique_ptr<Impl> impl_;
    };

}  // namespace nebula::io
